import os
import yaml
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    ServiceContext,
    load_index_from_storage
)
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core.query_engine import RetrieverQueryEngine


def read_api_from_config():
    config_filename = "config/config2.yaml"
    with open(config_filename, "r") as f:
        config = yaml.safe_load(f)
    return config["rag_llm"]["api_key"], config["rag_llm"]["base_url"]


def get_service_context():
    api_key, api_base = read_api_from_config()
    embed_model = OpenAIEmbedding(api_base=api_base, api_key=api_key)
    llm = OpenAI(model="gpt-3.5-turbo",api_base=api_base, api_key=api_key)
    service_context = ServiceContext.from_defaults(
        embed_model=embed_model, llm=llm)
    return service_context


def init_query_engine(service_context: ServiceContext):

    assert "rag" in os.listdir("."), "This script should be executed within root directory of the project."

    # check if data indexes already exists
    if not os.path.exists("rag/storage"):
        print("Indexing rust tutorials, please wait...")
        
        # load document
        documents = SimpleDirectoryReader(input_dir="rag/dataFiles").load_data(show_progress=True)
        print("Length of documents: ", len(documents))

        node_parser = SimpleNodeParser.from_defaults(chunk_size=1024)
        base_nodes = node_parser.get_nodes_from_documents(documents=documents)
        print("Base nodes length: ", len(base_nodes))
        print(base_nodes[0])
        
        # creating index
        index = VectorStoreIndex(nodes=base_nodes, service_context=service_context)
        # store index
        index.storage_context.persist(persist_dir="rag/storage")
    else:
        print("Loading rust tutorial indexing...")
        # load existing index
        storage_context = StorageContext.from_defaults(persist_dir="rag/storage")
        index = load_index_from_storage(storage_context=storage_context, service_context=service_context)

    # create retriever
    retriever = index.as_retriever(similarity_top_k=3)
    # query retriever engine
    query_engine = RetrieverQueryEngine.from_args(
        retriever=retriever,
        service_context=service_context
    )

    return query_engine


def query(content: str, query_engine: RetrieverQueryEngine):
    # test response
    response = query_engine.query(content)
    return response

QUERY_ENGINE = init_query_engine(get_service_context())

#! Unit test
if __name__ == "__main__":
    response = query("How to create a function in rust?", QUERY_ENGINE)
    print(response)
